# -*- coding: utf-8 -*-
"""
@Time    : 2025/3/10 10:36 
@Author  : ZhangShenao 
@File    : 3. 调用本地模型.py 
@Desc    : 
"""

# 导入transformers依赖库
# AutoModelForCausalLM代表模型
# AutoTokenizer代表分词器
# pipeline是Hugging Face提供的一个高级工具,用于简化模型的调用

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# 指定模型文件保存路径
cache_dir = "./my_model/uer/gpt2-chinese-cluecorpussmall/models--uer--gpt2-chinese-cluecorpussmall/snapshots/c2c0249d8a2731f269414cc3b22dff021f8e07a3"

# 设置推理设备类型
# 如果有有CUDA环境,则使用CUDA进行推理;否则采用CPU
device = "cuda" if torch.cuda.is_available() else "cpu"

# 加载模型和分词器
model = AutoModelForCausalLM.from_pretrained(cache_dir)
tokenizer = AutoTokenizer.from_pretrained(cache_dir)

# 使用加载的模型和分词器创建生成文本的pipeline
generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)

# 生成文本
output = generator(
    "请帮我讲一个冷笑话：",
    max_length=100,
    num_return_sequences=1,
    truncation=True,
    temperature=0.7,
    top_k=50,
    top_p=0.9,
    clean_up_tokenization_spaces=True
)
print(output)
